import unittest

import cv2
import numpy as np
from ocr_process.cls.clas_pipeline import ClasPipeline
from ocr_process.cls.image_orientation_pipeline import OrientationPipeline
from ocr_process.det.det_pipline import DetPipeline
from ocr_process.det.predict_det import TextDetector
from ocr_process.rec.predict_system import TextSystem
from ocr_process.rec.text_pipline import TextPipeline
from paddlenlp.transformers import ErnieLayoutTokenizer


class ClasTestCase(unittest.TestCase):

    def test_clas_1(self):
        a = np.array([1, 1])
        b = np.array([3, 2])
        print(b - a)

    def test_uie(self):
        tokenizer = ErnieLayoutTokenizer.from_pretrained('uie-x-base')
        print(tokenizer)

    def test_text_pipeline(self):
        text_pipeline = TextPipeline(det_model_path=r"D:\ocr\onnx\deploy\text_system\det\inference.onnx",
                                     rec_model_path=r"D:\ocr\onnx\deploy\text_system\rec\inference.onnx")
        image_data = cv2.imread("1.jpg", cv2.IMREAD_COLOR)
        input_data = {"original_image": image_data, "image": image_data.copy()}
        text_tesult = text_pipeline(input_data)
        print(text_tesult)

    def test_text_system(self):
        text_system = TextSystem(onnx_home=r"D:\ocr\onnx\deploy\text_system")
        image_data = cv2.imread("1.jpg", cv2.IMREAD_COLOR)
        text_tesult = text_system(image_data)
        print(text_tesult)

    def test_det_pipeline(self):
        det = DetPipeline(model_path=r"D:\ocr\onnx\deploy\text_system\det\inference.onnx")
        image_data = cv2.imread("1.jpg", cv2.IMREAD_COLOR)
        input_data = {"original_image": image_data, "image": image_data.copy()}
        result = det(input_data)
        print(result['dt_boxes'].shape)
        print(result['elapse'])

    def test_text_det_pipeline(self):
        text_detector = TextDetector(onnx_home=r"D:\ocr\onnx\deploy\text_system")
        image_data = cv2.imread("1.jpg", cv2.IMREAD_COLOR)
        dt_boxes, elapse = text_detector(image_data)
        # (29, 4, 2)
        print(dt_boxes.shape)
        print(elapse)

    def test_clas_pipeline(self):
        clas = ClasPipeline(model_path=r"D:\ocr\onnx\deploy\clas\inference.onnx")
        image_data = cv2.imread("1.jpg", cv2.IMREAD_COLOR)
        input_data = {"original_image": image_data, "image": image_data.copy()}
        result = clas(input_data)
        print(result)

    def test_orientation_pipeline(self):
        orientationPipeline = OrientationPipeline(model_path=r"D:\ocr\onnx\deploy\text_image\inference.onnx")
        image_data = cv2.imread("1.jpg", cv2.IMREAD_COLOR)
        input_data = {"original_image": image_data, "image": image_data.copy()}
        result = orientationPipeline(input_data)
        print(result)
